﻿// Copyright (c) 2025 liucl
// Inject is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan
// PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

#include "injector.hpp"

#include <psapi.h>
#include <tchar.h>
#include <windows.h>

#include <boost/scope_exit.hpp>

std::wstring absolutePath(const std::wstring &path) {
  wchar_t buffer[MAX_PATH] = {0};
  auto length = GetFullPathNameW(path.c_str(), MAX_PATH, buffer, NULL);
  return std::wstring(buffer);
}

bool Injector::Inject(DWORD pid, const std::wstring &src) {
  auto module = absolutePath(src);
  // 打开目标进程
  auto hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
  if (hProcess == nullptr) {
    return false;
  }
  BOOST_SCOPE_EXIT(hProcess) { CloseHandle(hProcess); }
  BOOST_SCOPE_EXIT_END;

  // 获取 Kernel32.dll 模块句柄
  auto hModule = GetModuleHandleW(_T("Kernel32.dll"));
  if (hModule == nullptr) {
    return false;
  }

  // 获取 LoadLibraryW 函数地址
  auto pfnLoadLibraryW = GetProcAddress(hModule, "LoadLibraryW");
  if (pfnLoadLibraryW == nullptr) {
    return false;
  }

  // 在目标进程中分配内存以存储 DLL 路径
  SIZE_T dwSize = (module.length() + 1) * sizeof(wchar_t);
  LPVOID lpBaseAddress = VirtualAllocEx(
      hProcess, NULL, dwSize, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
  if (lpBaseAddress == nullptr) {
    return false;
  }

  BOOST_SCOPE_EXIT(lpBaseAddress, hProcess) {
    VirtualFreeEx(hProcess, lpBaseAddress, 0, MEM_RELEASE);
  }
  BOOST_SCOPE_EXIT_END;

  // 将 DLL 路径写入目标进程的内存
  if (!WriteProcessMemory(hProcess, lpBaseAddress, module.c_str(), dwSize,
                          NULL)) {
    return false;
  }

  // 创建远程线程来执行 LoadLibraryW
  HANDLE hThread = CreateRemoteThread(hProcess, NULL, 0,
                                      (LPTHREAD_START_ROUTINE)pfnLoadLibraryW,
                                      lpBaseAddress, 0, NULL);
  if (hThread == nullptr) {
    return false;
  }
  BOOST_SCOPE_EXIT(hThread) { CloseHandle(hThread); }
  BOOST_SCOPE_EXIT_END;

  // 等待线程执行完成
  WaitForSingleObject(hThread, INFINITE);
  return true;
}
// 辅助函数：获取目标进程中已加载模块的句柄
static HMODULE GetRemoteModuleHandle(HANDLE hProcess,
                                     const wchar_t *moduleName) {
  HMODULE hModules[1024];
  DWORD cbNeeded;

  // 枚举目标进程中的所有模块
  if (EnumProcessModules(hProcess, hModules, sizeof(hModules), &cbNeeded)) {
    for (unsigned int i = 0; i < (cbNeeded / sizeof(HMODULE)); i++) {
      wchar_t szModName[MAX_PATH];

      // 获取模块的完整路径
      if (GetModuleFileNameExW(hProcess, hModules[i], szModName,
                               sizeof(szModName) / sizeof(wchar_t))) {
        // 比较模块名称
        if (wcsstr(szModName, moduleName)) {
          return hModules[i];
        }
      }
    }
  }
  return nullptr;
}

bool Injector::UnInject(DWORD pid, const std::wstring &src) {
  auto module = absolutePath(src);
  // 打开目标进程
  auto hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
  if (hProcess == nullptr) {
    return false;
  }
  BOOST_SCOPE_EXIT(hProcess) { CloseHandle(hProcess); }
  BOOST_SCOPE_EXIT_END;

  HMODULE hModule = GetRemoteModuleHandle(hProcess, module.c_str());
  if (hModule == nullptr) {
    return false;
  }

  // 获取 Kernel32.dll 模块句柄
  auto hKernel32 = GetModuleHandleW(_T("Kernel32.dll"));
  if (hKernel32 == nullptr) {
    return false;
  }

  // 获取 FreeLibrary 函数地址
  auto pfnFreeLibrary = GetProcAddress(hKernel32, "FreeLibrary");
  if (pfnFreeLibrary == nullptr) {
    return false;
  }

  // 创建远程线程来执行 FreeLibrary
  HANDLE hThread = CreateRemoteThread(hProcess, NULL, 0,
                                      (LPTHREAD_START_ROUTINE)pfnFreeLibrary,
                                      hModule, 0, NULL);
  if (hThread == nullptr) {
    return false;
  }
  BOOST_SCOPE_EXIT(hThread) { CloseHandle(hThread); }
  BOOST_SCOPE_EXIT_END;

  // 等待线程执行完成
  WaitForSingleObject(hThread, INFINITE);
  return true;
}